"""
ComfyUI 主流 LLM 節點集合
支援 Llama、Mistral、Phi、Qwen、Gemma、Yi 等主流開源模型
"""

from .llama_node import LlamaLLMNode
from .mistral_node import MistralLLMNode
from .phi_node import PhiLLMNode
from .qwen_node import QwenLLMNode
from .gemma_node import GemmaLLMNode
from .yi_node import YiLLMNode

# ComfyUI 節點註冊
NODE_CLASS_MAPPINGS = {
    "LlamaLLMNode": LlamaLLMNode,
    "MistralLLMNode": MistralLLMNode,
    "PhiLLMNode": PhiLLMNode,
    "QwenLLMNode": QwenLLMNode,
    "GemmaLLMNode": GemmaLLMNode,
    "YiLLMNode": YiLLMNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LlamaLLMNode": "🦙 Llama LLM Node",
    "MistralLLMNode": "🌪️ Mistral LLM Node", 
    "PhiLLMNode": "⚡ Phi LLM Node",
    "QwenLLMNode": "🇨🇳 Qwen LLM Node",
    "GemmaLLMNode": "💎 Gemma LLM Node",
    "YiLLMNode": "🔥 Yi LLM Node",
}

# 版本信息
__version__ = "1.0.0"
__author__ = "ComfyUI-MyLLMNode"
__description__ = "Universal LLM nodes for mainstream open-source models"

# 支援的模型列表
SUPPORTED_MODELS = {
    "Llama": ["Llama-2-7B", "Llama-2-13B", "Code-Llama", "Alpaca", "Vicuna"],
    "Mistral": ["Mistral-7B-Instruct", "Mixtral-8x7B"],
    "Phi": ["Phi-3-Mini", "Phi-3-Medium", "Phi-4-Mini"],
    "Qwen": ["Qwen-1.5-7B", "Qwen-2-7B", "Qwen-2.5-7B"],
    "Gemma": ["Gemma-2B", "Gemma-7B", "Gemma-2-9B"],
    "Yi": ["Yi-6B", "Yi-9B", "Yi-34B"]
}

print(f"Loaded ComfyUI-MyLLMNode v{__version__}")
print(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")